# pip install torch torchvision pillow pandas pymupdf
import random, re
from pathlib import Path
from typing import Optional
import pandas as pd
from PIL import Image
import torch
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T

try:
    import fitz  # PyMuPDF
except Exception:
    fitz = None

# -------- CLIP 标准归一化 --------
CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD  = (0.26862954, 0.26130258, 0.27577711)

def build_image_train_transform(img_size=224):
    return T.Compose([
        T.RandomResizedCrop(img_size, scale=(0.8, 1.0), antialias=True),
        T.RandomHorizontalFlip(p=0.5),
        T.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.02),
        T.RandomGrayscale(p=0.05),
        T.RandomPerspective(distortion_scale=0.2, p=0.1),
        T.ToTensor(),
        T.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
    ])

def build_image_eval_transform(img_size=224):
    return T.Compose([
        T.Resize(int(img_size * 1.14), antialias=True),
        T.CenterCrop(img_size),
        T.ToTensor(),
        T.Normalize(mean=CLIP_MEAN, std=CLIP_STD),
    ])

# -------- Tiny 文本增强 --------
_punct_re = re.compile(r"[^\w\s]")

def text_dropout(words, p=0.05):
    return [w for w in words if random.random() > p or len(words) <= 4]

def text_shuffle_minor(words, p=0.05):
    ws = words[:]
    for i in range(len(ws)-1):
        if random.random() < p:
            ws[i], ws[i+1] = ws[i+1], ws[i]
    return ws

def augment_text(text: str, do_prob=0.3):
    if random.random() > do_prob:
        return text.strip()
    t = text.strip()
    t = _punct_re.sub(lambda m: m.group(0) if random.random() < 0.7 else " ", t)
    words = t.split()
    words = text_dropout(words, p=0.05)
    words = text_shuffle_minor(words, p=0.05)
    return " ".join(words)

# -------- PDF 渲染 --------
def _render_pdf_page_to_pil(pdf_path: Path, page_index=0, dpi=196) -> Image.Image:
    if fitz is None:
        raise RuntimeError("请先安装 pymupdf: pip install pymupdf")
    doc = fitz.open(pdf_path)
    page_index = max(0, min(page_index, len(doc)-1))
    page = doc[page_index]
    zoom = dpi / 72.0
    mat = fitz.Matrix(zoom, zoom)
    pix = page.get_pixmap(matrix=mat, alpha=False)
    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
    doc.close()
    return img

def load_image_any(path: Path, train=True, deterministic_page: Optional[int]=None, pdf_dpi=196):
    suffix = path.suffix.lower()
    if suffix == ".pdf":
        with fitz.open(path) as doc:
            n = len(doc)
        page_idx = random.randrange(n) if train else (deterministic_page or 0)
        return _render_pdf_page_to_pil(path, page_idx, dpi=pdf_dpi)
    else:
        return Image.open(path).convert("RGB")

# -------- Dataset --------
class PairsDataset(Dataset):
    """
    CSV 有表头: image_path, caption
    支持 pdf/png/jpg 等
    """
    def __init__(self, csv_path, image_root="", train=True,
                 img_size=224, text_aug_prob=0.3, pdf_render_dpi=196, deterministic_page=None):
        self.df = pd.read_csv(csv_path)
        self.root = Path(image_root) if image_root else None
        self.train = train
        self.text_aug_prob = text_aug_prob if train else 0.0
        self.pdf_dpi = pdf_render_dpi
        self.det_page = deterministic_page
        self.img_tf = build_image_train_transform(img_size) if train else build_image_eval_transform(img_size)

    def __len__(self): return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_fp = Path(row["Source"])
        if self.root and not img_fp.is_absolute():
            img_fp = (self.root / img_fp).resolve()
        image = load_image_any(img_fp, train=self.train, deterministic_page=self.det_page, pdf_dpi=self.pdf_dpi)
        image = image.convert("RGB")
        image = self.img_tf(image)
        caption = str(row["Caption"])

        if self.train:
            caption = augment_text(caption, do_prob=self.text_aug_prob)
        return image, caption, idx

def build_loader(csv_path, image_root="", batch_size=64, num_workers=4,
                 train=True, img_size=224, text_aug_prob=0.3, pdf_render_dpi=196, deterministic_page=None):
    ds = PairsDataset(csv_path, image_root, train, img_size, text_aug_prob, pdf_render_dpi, deterministic_page)
    dl = DataLoader(ds, batch_size=batch_size, shuffle=train,
                    num_workers=num_workers, pin_memory=True, drop_last=train)
    return dl, ds

# ---------------- 使用示例 ----------------
# 验证 (固定 PDF 第一页):
# val_loader, _ = build_loader("train.csv", train=False, deterministic_page=0)
